#ifndef TINYNDARRAY_NDARRAY_H
#define TINYNDARRAY_NDARRAY_H

namespace tinyndarray {

class ConstElementProxy;
class ElementProxy;
class NdArray;
class NdArrayView;
class Substance {
public:
    Dtype dtype;
    size_t size;
    Shape shape;
    std::shared_ptr<void> data;
    bool owns_data;
    std::vector<size_t> strides;

    // 拥有数据的构造函数
    Substance(Dtype dtype_, size_t size_, const Shape& shape_)
        : dtype(dtype_),
          size(size_),
          shape(shape_),
          data(new char[size_ * get_type_size(dtype_)], std::default_delete<char[]>()),
          owns_data(true) {
        strides = compute_strides(shape, get_type_size(dtype_));
    }

    // 使用外部数据的构造函数
    Substance(Dtype dtype_, size_t size_, const Shape& shape_, void* external_data)
        : dtype(dtype_),
          size(size_),
          shape(shape_),
          data(external_data, [](void*) {}),  // 空删除器
          owns_data(false) {
        strides = compute_strides(shape, get_type_size(dtype_));
    }
};

// 通用的嵌套初始化列表模板
template <typename T, std::size_t D>
struct NumListHelper {
    using type = std::initializer_list<typename NumListHelper<T, D - 1>::type>;
};

template <typename T>
struct NumListHelper<T, 0> {
    using type = std::initializer_list<T>;
};

template <typename T, std::size_t D>
using NumList = typename NumListHelper<T, D>::type;

template <typename NList>
std::list<int> CheckNListShapeImpl(const NList& init_list) {
    if (init_list.size() == 0) {
        return {};
    }
    // Check all children have same shape
    auto itr = init_list.begin();
    auto shape = CheckNListShapeImpl(*itr);
    for (size_t i = 0; i < init_list.size(); i++, itr++) {
        if (shape != CheckNListShapeImpl(*itr)) {
            throw std::runtime_error("Initializing shape is invalid");
        }
    }
    // Return total shape of children
    shape.push_front(static_cast<int>(init_list.size()));
    return shape;
}

template <typename NList>
Shape CheckNListShape(const NList& init_list) {
    // Check and get the shape of nested initializer.
    const std::list<int>& shape = CheckNListShapeImpl(init_list);
    // Cast to vector
    return Shape(shape.begin(), shape.end());
}

#define COPY_ELEMES_TEMP(ctype)                                                        \
    template <>                                                                        \
    inline std::list<int> CheckNListShapeImpl(const NumList<ctype, 0>& init_list) {    \
        return {static_cast<int>(init_list.size())};                                   \
    }                                                                                  \
    template <typename NList>                                                          \
    inline void CopyNListElemsImpl(const NList& init_list, ctype*& data) {             \
        for (auto itr = init_list.begin(); itr != init_list.end(); itr++) {            \
            CopyNListElemsImpl(*itr, data);                                            \
        }                                                                              \
    }                                                                                  \
    template <>                                                                        \
    inline void CopyNListElemsImpl(const NumList<ctype, 0>& init_list, ctype*& data) { \
        for (auto&& v : init_list) {                                                   \
            *(data++) = v;                                                             \
        }                                                                              \
    }                                                                                  \
    template <typename NList>                                                          \
    inline void CopyNListElems(const NList& init_list, ctype* data) {                  \
        CopyNListElemsImpl(init_list, data);                                           \
    }

COPY_ELEMES_TEMP(float32_t)
COPY_ELEMES_TEMP(int)
#undef COPY_ELEMES_TEMP

// 主类声明
class NdArray {
public:
    template <bool Const>
    class IterBase;
    using Iter = IterBase<false>;
    using ConstIter = IterBase<true>;

    static std::mt19937 s_rand_engine;
    static std::atomic<uint32_t> s_rand_seed;
    static std::mutex s_rand_mutex;

    // 构造和析构
    NdArray();
    NdArray(Dtype dtype, const Shape& shape);
    NdArray(Dtype dtype, const Shape& shape, float32_t fill_value);
    NdArray(Dtype dtype, const Shape& shape, const void* external_data);
    // 新增构造函数：使用外部数据（无所有权）
    NdArray(Dtype dtype, const Shape& shape, void* external_data, bool own_data = false);
    NdArray(const NdArray&) = default;
    NdArray(NdArray&&) noexcept = default;
    ~NdArray() = default;
#define DECL_NDARRAY_INIT_LIST(ctype)     \
    NdArray(NumList<ctype, 0> init_list); \
    NdArray(NumList<ctype, 1> init_list); \
    NdArray(NumList<ctype, 2> init_list); \
    NdArray(NumList<ctype, 3> init_list); \
    NdArray(NumList<ctype, 4> init_list); \
    NdArray(NumList<ctype, 5> init_list); \
    NdArray(NumList<ctype, 6> init_list); \
    NdArray(NumList<ctype, 7> init_list); \
    NdArray(NumList<ctype, 8> init_list); \
    NdArray(NumList<ctype, 9> init_list);

    DECL_NDARRAY_INIT_LIST(float32_t)
    DECL_NDARRAY_INIT_LIST(int)
    NdArray astype(Dtype new_dtype) const;
    // 赋值运算符
    NdArray& operator=(const NdArray&) = default;
    NdArray& operator=(NdArray&&) = default;

    // 静态创建方法
    static NdArray Empty(Dtype dtype, const Shape& shape);
    static NdArray Zeros(Dtype dtype, const Shape& shape);
    static NdArray Ones(Dtype dtype, const Shape& shape);

    // 类型转换
    template <typename T>
    T as_scalar() const;

    void* data_ptr();
    const void* data_ptr() const;
    // 检查是否拥有数据
    bool owns_data() const;

    template <typename T>
    T* typed_data();

    template <typename T>
    const T* typed_data() const;

    std::string to_string() const;
    // 基本属性
    uintptr_t id() const;
    bool empty() const;
    size_t size() const;
    const Shape& shape() const;
    const int& shape(int axis) const;
    int itemsize() const;
    size_t ndim() const;
    Dtype dtype() const;

    // 数据操作
    template <typename T>
    void fill(T value);
    NdArray copy() const;
    void resize(const Shape& new_shape);

    // 迭代器
    Iter begin();
    Iter end();
    ConstIter begin() const;
    ConstIter end() const;

    // 元素访问
    ElementProxy operator[](size_t index);
    ElementProxy operator[](const Index& index);
    ConstElementProxy operator[](size_t index) const;
    ConstElementProxy operator[](const Index& index) const;

    template <typename... I>
    ElementProxy operator()(I... indices);

    // 形状操作
    NdArray reshape(const Shape& new_shape) const;
    NdArray flatten() const;
    NdArray ravel() const;

    NdArrayView slice(const SliceIndex& slice_index) const;
    template <typename... I>
    NdArrayView slice(std::initializer_list<I>... slice_index) const;
    NdArray slice_copy(const SliceIndex& slice_index) const;
    template <typename... I>
    NdArray slice_copy(std::initializer_list<I>... slice_index) const;

    NdArrayView operator()(const SliceIndex& slice_index) const;
    template <typename... I>
    NdArrayView operator()(std::initializer_list<I>... slice_index) const;
    // 点积和叉积
    NdArray dot(const NdArray& other) const;
    NdArray cross(const NdArray& other) const;

    // 归约操作
    NdArray sum(const Axis& axes = {}, bool keepdims = false) const;
    NdArray mean(const Axis& axes = {}, bool keepdims = false) const;
    NdArray min(const Axis& axes = {}, bool keepdims = false) const;
    NdArray max(const Axis& axes = {}, bool keepdims = false) const;

    // 文件 I/O 方法
    void Save(const std::string& file_path) const;

    // 从 cnpy 数组加载
    static NdArray FromCnpyArray(const void* cnpy_array);

    // 获取 cnpy 类型
    const char* get_cnpy_type() const;

    // 获取 cnpy 类型大小
    size_t get_cnpy_type_size() const;

    // 获取 cnpy 类型字符
    char get_cnpy_type_char() const;

    // 类型转换操作符
    operator float32_t() const;
    operator int() const;
    operator float64_t() const;

    const std::vector<size_t>& strides() const;

    // 随机数
    static void Seed();
    static void Seed(uint32_t seed);
    static NdArray Uniform(float16_t low, float16_t high, const Shape& shape);
    static NdArray Uniform(float32_t low, float32_t high, const Shape& shape);
    static NdArray Uniform(float64_t low, float64_t high, const Shape& shape);
    static NdArray Uniform(int low, int high, const Shape& shape);
    static NdArray Uniform(const Shape& shape);
    static NdArray Normal(float32_t loc, float32_t scale, const Shape& shape);
    static NdArray Normal(float64_t loc, float64_t scale, const Shape& shape);
    static NdArray Normal(const Shape& shape);

    // 辅助方法
    size_t calculate_offset(const Index& index) const;
    static size_t calculate_size(const Shape& shape);
    static std::mt19937 GetThreadSafeRandomEngine();

public:
    std::shared_ptr<Substance> m_sub;
};
NdArray Load(const std::string& file_path);

// 非成员函数声明
std::ostream& operator<<(std::ostream& os, const NdArray& arr);
std::ostream& operator<<(std::ostream& os, const Shape& shape);
NdArray operator+(const NdArray& lhs, const NdArray& rhs);
NdArray operator-(const NdArray& lhs, const NdArray& rhs);
NdArray operator*(const NdArray& lhs, const NdArray& rhs);
NdArray operator/(const NdArray& lhs, const NdArray& rhs);
NdArray operator==(const NdArray& lhs, const NdArray& rhs);
NdArray operator<=(const NdArray& lhs, const NdArray& rhs);
NdArray operator>=(const NdArray& lhs, const NdArray& rhs);
NdArray operator<(const NdArray& lhs, const NdArray& rhs);
NdArray operator>(const NdArray& lhs, const NdArray& rhs);
template <typename T>
NdArray operator+(const NdArray& lhs, const T& rhs);
extern template NdArray operator+(const NdArray& lhs, const float32_t& rhs);

template <typename T>
NdArray operator*(const NdArray& lhs, const T& rhs);
extern template NdArray operator*(const NdArray& lhs, const float32_t& rhs);

template <typename T>
NdArray operator/(const NdArray& lhs, const T& rhs);
extern template NdArray operator/(const NdArray& lhs, const float32_t& rhs);
NdArray operator-(const NdArray& x);

NdArray Exp(const NdArray& x);

// 形状操作函数
NdArray Reshape(const NdArray& x, const Shape& shape);
NdArray Squeeze(const NdArray& x, const Axis& axes = {});
NdArray ExpandDims(const NdArray& x, int axis);
NdArray Transpose(const NdArray& x, const Axis& perm);
NdArray SwapAxes(const NdArray& x, int axis1, int axis2);
NdArray BroadcastTo(const NdArray& x, const Shape& new_shape);

// 组合操作
NdArray Stack(const std::vector<NdArray>& xs, int axis);
NdArray Concatenate(const std::vector<NdArray>& xs, int axis);
std::vector<NdArray> Split(const NdArray& x, int n_section, int axis);
std::vector<NdArray> Split(const NdArray& x, const Index& idxs, int axis);
std::vector<NdArray> Separate(const NdArray& x, int axis);
}  // namespace tinyndarray

#endif